#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Licensed Materials - Property of tenxcloud.com
# (C) Copyright 2020 TenxCloud. All Rights Reserved.
# 2020-06-08 @author lizhen

import tensorflow as tf

from pathlib import Path
from datetime import datetime
import os


class TenxLogger(tf.keras.callbacks.Callback):

    def __init__(self):
        super(TenxLogger, self).__init__()

    def on_epoch_end(self, epoch, logs=None):
        if logs is None:
            return
        now = datetime.now()
        part_one = now.strftime("%Y-%m-%d %H:%M:%S,%f")
        part_one = part_one[:-3]
        part_two = now.strftime("%H:%M:%S")
        print("\n{:s} INFO 	{:s} Te={:d} Loss={:.3f} | AccT={:.3f}%\n".format(
            part_one, part_two, epoch + 1, float(logs.get('loss', 0)) * 100, float(logs.get('accuracy', 0)) * 100))

    # def on_train_batch_end(self, batch, logs=None):
    #     if logs is None:
    #         return
    #     print("batch", batch, logs)


parent = Path(os.path.dirname(os.path.abspath(__file__))).parent
model_set_dir = os.path.join(parent, 'modelsets')
mnist_npz = os.path.join(parent, 'datasets', 'mnist.npz')

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data(path=mnist_npz)
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10)
])

predictions = model(x_train[:1]).numpy()
tf.nn.softmax(predictions).numpy()
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
loss_fn(y_train[:1], predictions).numpy()

model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])

if os.path.exists(os.path.join(model_set_dir, 'checkpoint')):
    latest = tf.train.latest_checkpoint(model_set_dir)
    model.load_weights(latest)

checkpoint_path = os.path.join(model_set_dir, 'cp-{epoch:04d}.ckpt')

cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=0)
log_callback = TenxLogger()

model.fit(x_train, y_train, epochs=5, callbacks=[cp_callback, log_callback])
model.evaluate(x_test, y_test, verbose=0)

probability_model = tf.keras.Sequential([
    model,
    tf.keras.layers.Softmax()
])

probability_model(x_test[:5])

# model.save(os.path.join(model_set_dir, 'mnist.h5'))
model.save(os.path.join(model_set_dir, 'saved_model', "{model_version:d}".format(model_version=1)))
